/** 
 * @file test.cpp
 * @author Lin Lin
 * @date 2006-12-27
 */


#include	<iostream>
#include	<new>
#include	<cassert>
#include	<complex>
#include	<cmath>

#include	<utility>
#include	<vector>
#include	<algorithm>

#include	<mpi.h>
#include	<fftw3.h>

using namespace std;

typedef unsigned int UINT;



/** 
 * @defgroup GlobVar global variables
 *   所有的全局变量的说明。
 *            
 */

/**
 * @addtogroup GlobVar
 * 时间和时间步长
 * @{
 */
double t;		/**< 时间 */
double dt;		/**< 时间步长 */
/** @}*/

/**
 * @addtogroup GlobVar
 * 表征全局和局部的数组大小的量
 *
 * @{
 */
const int DIM = 2;          /**< 区域的维数 */
int n[DIM];	        /**< 整个区域的大小 */
int local_n[DIM];	/**< 单个节点上区域的大小 */
int total_local_size;    /**< 一个节点上的数组的大小  */

int N[DIM]; /**< 区域在各个方向都 1.5 倍以后的大小 */
int local_N[DIM];
int local_start[DIM];    /**< 单个节点上数组的起始指标 */
int nproc;	/**< 总的节点个数 */
int rank;	/**< 本节点的秩 */
/** @}*/

double wtime;    /**< 墙钟 */
double h[DIM];   /**< 空间步长 */ 

/**
 * @addtogroup GlobVar
 * 预定义的MPI数据类型
 *
 * @{
 */
MPI_Datatype MPI_FFTW_COMPLEX;
MPI_Datatype MPI_LOCAL_COLUMN;
/** @}*/

/**
 * @addtogroup GlobVar
 * 用到的FFTW的plan: Guru Interface
 *
 * @{
 */
fftw_plan fftwp_r2c;
fftw_plan fftwp_c2c;
fftw_plan fftwp_c2r;
/** @}*/

/**
 * @addtogroup GlobVar
 * 各种常数
 *
 * @{
 */

const double PI = 3.14159265358979323844; 	
double normalize_scale;		/**< 规一化常数*/

/** @}*/

void err_handler(MPI_Comm * comm, int * err_code, ...)
{
  int resultlen;
  char err_string[MPI_MAX_ERROR_STRING];
  MPI_Error_string(*err_code, err_string, &resultlen);
  printf("MPI exception caughted: error code %d\n", *err_code);
  printf("\terror: %s\n", err_string);
  cout << "OK" << endl;
  MPI_Abort(*comm, *err_code);
}

template<typename Tp>
class DenseMatrix
{
	private:
		UINT M_m;
		UINT M_n;
		vector<Tp> M_vec_elem;	// row first.
	
	public:
	

		/////////////////////////////////////////////////
		/// Constructors and Destructors
		/////////////////////////////////////////////////
		DenseMatrix(const UINT m, const UINT n)
		  	:M_m(m), M_n(n)
		{
		  M_vec_elem.resize(m*n);
		}

		~DenseMatrix()
		{}

		DenseMatrix&
	    operator = (DenseMatrix& dnm_in)
		{
		  if( M_m != dnm_in.m() || M_n != dnm_in.n())
			M_vec_elem.resize(dnm_in.m() * dnm_in.n());
		  copy(dnm_in.M_vec_elem.begin(), dnm_in.M_vec_elem.end(),
			   M_vec_elem.begin());
		}


		/////////////////////////////////////////////////
		/// Non-modifying operators
		/////////////////////////////////////////////////
		UINT
		m() const
		{
		  return M_m;
		}
	
		UINT
		n() const
		{
		  return M_n;
		}

		const Tp&
		operator() (const UINT i, const UINT j) const
		{
		  return M_vec_elem[i * M_n + j];
		}

		void print()
		{
		  for(UINT i = 0; i < M_m; i++)
		  {
			for(UINT j = 0; j < M_n; j++)
			{
			  //cout << M_vec_elem[j + i*M_n] << " ";
			  cout << (*this)(i, j) << " ";
			}
			cout << endl;
		  }
		  cout << endl;
		}

		/////////////////////////////////////////////////
		/// Modifying operators
		/////////////////////////////////////////////////
		Tp&
		operator() (const UINT i, const UINT j)
		{
		  return M_vec_elem[i * M_n + j];
		}
		
		DenseMatrix&
		operator *= (const Tp& factor)
		{
		  for(iterator iter = begin(); iter != end(); iter++)
			*iter *= factor;
		}
		
		/** 
		 * @brief 
		 *
		 * Safety checking has not performed.
		 *
		 * @param factor 
		 * 
		 * @return 
		 */
		DenseMatrix&
		operator /= (const Tp& factor)
		{
		  for(iterator iter = begin(); iter != end(); iter++)
			*iter /= factor;
		}
		
		DenseMatrix&
		add_scaled (const Tp& factor, const DenseMatrix& dnm_in)
		{
		  for(iterator iter = begin(), iter_in = dnm_in.begin();
			  iter != end(); 
			  iter++, iter_in++)
			*iter += factor * (*iter_in);
		}


		/////////////////////////////////////////////////
		///	Iterator Functions
		/////////////////////////////////////////////////
		typedef class vector<Tp>::const_iterator const_iterator;
		typedef class vector<Tp>::iterator iterator;

		const_iterator
		begin() const
		{
		  return M_vec_elem.begin();
		}

		const_iterator
		end() const
		{
		  return M_vec_elem.end();
		}
		
		const_iterator
		begin(const UINT row) const
		{
		  return M_vec_elem.begin() + row * M_n;
		}

		const_iterator
		end(const UINT row) const
		{
		  return M_vec_elem.end() + row * (M_n+1);
		}

		iterator
		begin()
		{
		  return M_vec_elem.begin();
		}

		iterator
		end()
		{
		  return M_vec_elem.end();
		}
		
		iterator
		begin(const UINT row)
		{
		  return M_vec_elem.begin() + row * M_n;
		}

		iterator
		end(const UINT row)
		{
		  return M_vec_elem.end() + row * (M_n+1);
		}

};




int main(int argc, char* argv[])
{
  MPI_Status status;
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &nproc);
  MPI_Errhandler errhdl;
  MPI_Errhandler_create(&err_handler, &errhdl);
  MPI_Errhandler_set(MPI_COMM_WORLD, errhdl);
 
  /**
   * Initialize Constants
   *
   */
  N[0] = 8;
  N[1] = 8;
  local_N[0] = N[0]/nproc;
  local_N[1] = N[1];

  /**
   * Initialize Datatypes
   *
   */

//  MPI_Type_contiguous(2, MPI_DOUBLE, &MPI_FFTW_COMPLEX);
//  MPI_Type_commit(&MPI_FFTW_COMPLEX);
//  MPI_Type_vector(local_m, 1, global_n, MPI_FFTW_COMPLEX, &MPI_LOCAL_COLUMN);
//  MPI_Type_commit(&MPI_LOCAL_COLUMN);

  cout << "Rank: " << rank << ", preparing FFTW plans." << endl;
  /**
   * Initialize FFTW plans with guru interface
   *
   */
  DenseMatrix<double>	*pdnm_in = new DenseMatrix<double>(local_N[0], local_N[1]);
  DenseMatrix<complex<double> > *pdnm_out = new DenseMatrix<complex<double> >
											(local_N[0], local_N[1]/2 + 1);
  fftw_iodim dims_r2c[1];
  fftw_iodim howmany_dims_r2c[1];
  dims_r2c[0].n = local_N[1];
  dims_r2c[0].is = 1;
  dims_r2c[0].os = 1;
  howmany_dims_r2c[0].n = local_N[0] ;
  howmany_dims_r2c[0].is = local_N[1];
  howmany_dims_r2c[0].os = local_N[1]/2 + 1;
  
  fftwp_r2c = fftw_plan_guru_dft_r2c(1, dims_r2c, 1, howmany_dims_r2c, 
	 		    &(*pdnm_in)(0,0), reinterpret_cast<fftw_complex*>(&(*pdnm_out)(0,0)), 
	    	 	FFTW_MEASURE);
  delete pdnm_in, pdnm_out;
  
  DenseMatrix<double>	dnm_in(local_N[0], local_N[1]);
  DenseMatrix<complex<double> > dnm_out(local_N[0], local_N[1]/2 + 1);
  for( DenseMatrix<double>::iterator iter = dnm_in.begin();
	  							     iter != dnm_in.end();
									 iter++)
  {
	*iter = distance(dnm_in.begin(), iter);
  }
  dnm_in.print();
 

  cout << "Rank: " << rank << ", executing FFTW plans." << endl;
  fftw_execute_dft_r2c(fftwp_r2c, &dnm_in(0,0), 
	 		 reinterpret_cast<fftw_complex*>(&dnm_out(0,0)) );

  cout << "Rank: " << rank << ", printing the output values" << endl;
  dnm_out.print();

  fftw_destroy_plan(fftwp_r2c); 
  MPI_Errhandler_free(&errhdl);
//  MPI_Type_free(&MPI_FFTW_COMPLEX);
//  MPI_Type_free(&MPI_LOCAL_COLUMN); 
  MPI_Finalize();
  return 0;
  
}
